import asyncio
from typing import List, Optional
import motor.motor_asyncio
from pymongo import MongoClient
from config.settings import settings
from models.article import ArticleModel
import logging

logger = logging.getLogger(__name__)

class MongoDBManager:
    def __init__(self):
        self.client = None
        self.database = None
        self.collection = None
        
    async def connect(self):
        """建立数据库连接"""
        try:
            self.client = motor.motor_asyncio.AsyncIOMotorClient(settings.MONGODB_URL)
            self.database = self.client[settings.DATABASE_NAME]
            self.collection = self.database[settings.COLLECTION_NAME]
            
            # 测试连接
            await self.client.admin.command('ping')
            logger.info("MongoDB连接成功")
            
            # 创建索引
            await self.create_indexes()
            
        except Exception as e:
            logger.error(f"MongoDB连接失败: {e}")
            raise
    
    async def create_indexes(self):
        """创建数据库索引"""
        try:
            # 为URL创建唯一索引，避免重复抓取
            await self.collection.create_index("url", unique=True)
            # 为行业关键词和爬取时间创建索引
            await self.collection.create_index([
                ("industry_keyword", 1),
                ("crawl_time", -1)
            ])
            # 为域名创建索引
            await self.collection.create_index("source_domain")
            logger.info("数据库索引创建成功")
        except Exception as e:
            logger.error(f"创建索引失败: {e}")
    
    async def save_article(self, article: ArticleModel) -> bool:
        """保存文章到数据库"""
        try:
            article_dict = article.dict(by_alias=True, exclude_unset=True)
            result = await self.collection.insert_one(article_dict)
            logger.info(f"文章保存成功，ID: {result.inserted_id}")
            return True
        except Exception as e:
            if "duplicate key error" in str(e):
                logger.warning(f"文章URL已存在: {article.url}")
            else:
                logger.error(f"保存文章失败: {e}")
            return False
    
    async def get_articles_by_industry(self, industry_keyword: str, limit: int = 100) -> List[ArticleModel]:
        """根据行业关键词获取文章"""
        try:
            cursor = self.collection.find(
                {"industry_keyword": industry_keyword}
            ).limit(limit).sort("crawl_time", -1)
            
            articles = []
            async for doc in cursor:
                articles.append(ArticleModel(**doc))
            return articles
        except Exception as e:
            logger.error(f"查询文章失败: {e}")
            return []
    
    async def get_article_count_by_industry(self) -> dict:
        """获取各行业文章数量统计"""
        try:
            pipeline = [
                {"$group": {"_id": "$industry_keyword", "count": {"$sum": 1}}},
                {"$sort": {"count": -1}}
            ]
            cursor = self.collection.aggregate(pipeline)
            result = {}
            async for doc in cursor:
                result[doc["_id"]] = doc["count"]
            return result
        except Exception as e:
            logger.error(f"统计失败: {e}")
            return {}
    
    async def url_exists(self, url: str) -> bool:
        """检查URL是否已存在"""
        try:
            result = await self.collection.find_one({"url": url})
            return result is not None
        except Exception as e:
            logger.error(f"检查URL失败: {e}")
            return False
    
    async def close(self):
        """关闭数据库连接"""
        if self.client:
            self.client.close()
            logger.info("MongoDB连接已关闭")

# 全局数据库管理器实例
db_manager = MongoDBManager() 